import os
import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import argrelextrema
from copy import deepcopy

dt = 24 * 60 * 60
dt_times = 4500
T_times = 10
#save_path = os.path.dirname(__file__)
save_path = os.getcwd()

AU = 1.495978707e11
G = 6.67408e-11 / np.power(AU, 3) * 1e27 * dt * dt
M = {
    "Sun": 1.988e3,
    "Earth": 5.9722e-3,
    "Jupyter": 1.8986
}
ORBIT = {
    "Earth":{
        "a": 0.5 * (1.496e11 + 1.4958e11) / AU,
        "c": 0.5 * (1.496e11 - 1.4958e11) / AU
    },
    "Jupyter":{
        "a": (5.4588 + 4.9501) / 2,
        "c": (5.4588 - 4.9501) / 2
    }
}
PLANET = ["Earth", "Jupyter"]
POSITION = [
    ORBIT["Earth"]["a"] - ORBIT["Earth"]["c"],
    0,
    -(ORBIT["Jupyter"]["a"] + ORBIT["Jupyter"]["c"]),
    0
]
VELOCITY = [
    0,
    -np.sqrt(G * M["Sun"] * (2 * ORBIT["Earth"]["a"] - POSITION[0]) / (2 * ORBIT["Earth"]["a"] * POSITION[0])),
    0,
    np.sqrt(-G * M["Sun"] * (2 * ORBIT["Jupyter"]["a"] + POSITION[2]) / (2 * ORBIT["Jupyter"]["a"] * POSITION[2]))
]

def Simulation(position, velocity, t):
    x_e, y_e, x_j, y_j = deepcopy(position)
    vx_e, vy_e, vx_j, vy_j = deepcopy(velocity)
    r_ej = np.sqrt(np.square(x_e - x_j) + np.square(y_e - y_j))
    for i, x, y, vx, vy in zip([0, 1], [x_e, x_j], [y_e, y_j], [vx_e, vx_j], [vy_e, vy_j]):
        r = np.sqrt(np.square(x) + np.square(y))
        ax = -(x * np.power(r, -3) * M["Sun"] + np.power(-1, i + 2) * (x_e - x_j) * np.power(r_ej, -3) * M[PLANET[1 - i]]) * G
        ay = -(y * np.power(r, -3) * M["Sun"] + np.power(-1, i + 2) * (y_e - y_j) * np.power(r_ej, -3) * M[PLANET[1 - i]]) * G
        '''
        ax = -(x * np.power(r, -3) * M["Sun"]) * G
        ay = -(y * np.power(r, -3) * M["Sun"]) * G
        '''
        velocity[i * 2] = vx + ax * t
        velocity[i * 2 + 1] = vy + ay * t
        position[i * 2] = x + velocity[i * 2] * t
        position[i * 2 + 1] = y + velocity[i * 2 + 1] * t
    return position, velocity

if __name__ == "__main__":
    indexes = np.arange(0, dt_times * T_times, 1)
    positions = [POSITION]
    velocities = [VELOCITY]
    position = positions[-1]
    velocity = velocities[-1]
    for i, index in enumerate(indexes[:-1]):
        position, velocity = Simulation(positions[-1], velocities[-1], 1)
        positions.append(deepcopy(position))
        velocities.append(deepcopy(velocity))
    positions = np.array(positions).T
    velocities = np.array(velocities).T
    plt.subplot(311)
    plt.plot(positions[0], positions[1])
    plt.plot(positions[2], positions[3])
    plt.axis('equal')
    plt.subplot(312)
    r_e = np.sqrt(np.square(positions[0]) + np.square(positions[1]))
    r_e_max = np.array(argrelextrema(r_e, np.greater))[0]
    np.savetxt(os.path.join(save_path, "earth.txt"), r_e_max, fmt='%d')
    T_e = 2 * np.average(r_e_max[1:] - r_e_max[:-1])
    print(T_e, np.power(T_e, 2) / np.power(ORBIT["Earth"]["a"], 3))
    plt.plot(indexes[:int(T_e * 10)], r_e[:int(T_e * 10)])
    plt.grid()
    plt.subplot(313)
    r_j = np.sqrt(np.square(positions[2]) + np.square(positions[3]))
    r_j_max = np.array(argrelextrema(r_j, np.greater))[0]
    np.savetxt(os.path.join(save_path, "jupyter.txt"), r_j_max, fmt='%d')
    T_j = 2 * np.average(r_j_max[1:] - r_j_max[:-1])
    print(T_j, np.power(T_j, 2) / np.power(ORBIT["Jupyter"]["a"], 3))
    plt.plot(indexes, r_j)
    plt.grid()
    plt.savefig(os.path.join(save_path, "1.png"))
    plt.clf()
    plt.subplot(211)
    dSdt_e = (positions[0] * velocities[1] - positions[1] * velocities[0]) / 2
    plt.plot(indexes[:int(T_e * 10)], dSdt_e[:int(T_e * 10)])
    plt.grid()
    plt.subplot(212)
    dSdt_j = (positions[2] * velocities[3] - positions[3] * velocities[2]) / 2
    plt.plot(indexes, dSdt_j)
    plt.grid()
    plt.savefig(os.path.join(save_path, "2.png"))